import backtrader as bt
import math
from .DbOprate import DbOprate
import datetime

class SMACustom(bt.Indicator):
    lines = ('sma',)
    params = (('period', 2), ('weight', 1),)

    def next(self):
        if len(self) == 1:
            self.lines.sma[0] = self.data[0]
        else:
            m = self.params.weight
            n = self.params.period
            self.lines.sma[0] = (m * self.data[0] + (n - m) * self.lines.sma[-1]) / n

class TDXBSStrategy(bt.Strategy):
    params = (
        ('print_debug', True),   
        ('ema_period', 43),
        ('ma_period', 3),
        ('daily_output', True),
        ('log_to_file', True),
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.data_low = self.datas[0].low
        
        # 创建日志文件
        if self.params.log_to_file:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            self.log_file = open(f"tdx_strategy_log_{timestamp}.txt", "w", encoding="utf-8")
        
        # 技术指标计算 - 根据新公式
        self.p3 = bt.indicators.EMA(self.data_close, period=self.params.ema_period)  # P3:=EMA(CLOSE,43)
        self.x_7 = self.data_close  # X_7:=MA(CLOSE,1) 即收盘价本身
        self.sma_layer1 = SMACustom(self.x_7, period=2, weight=1)
        self.sma_layer2 = SMACustom(self.sma_layer1, period=2, weight=1)
        self.x_8 = SMACustom(self.sma_layer2, period=2, weight=1)  # X_8:=SMA(SMA(SMA(X_7,2,1),2,1),2,1)
        self.x_9 = bt.indicators.SimpleMovingAverage(self.x_8, period=self.params.ma_period)  # X_9:=MA(X_8,3)
        
        # 条件计算
        self.hl = self.x_8 >= self.x_9  # HL := X_8 >= X_9
        self.ll = self.x_8 <= self.x_9  # LL := X_8 <= X_9
        self.zzd = bt.And(self.hl, self.ll(-1))  # ZZD := HL AND REF(LL, 1)
        self.ll_count = bt.indicators.SumN(self.ll, period=10)  # COUNT(LL, 10)
        self.lxyz = (self.ll_count == 9)  # LXYZ := COUNT(LL, 10) = 9
        self.jichu_b = bt.And(self.zzd, self.lxyz, (self.data_close > self.p3))  # 基础B := ZZD AND LXYZ AND C>P3
        
        # 基础S - CROSS(X_9,L) 
        self.jichu_s = bt.indicators.CrossUp(self.x_9, self.data_low)  # 基础S := CROSS(X_9,L)

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.date(0)
        log_text = f'{dt.isoformat()}, {txt}'
        print(log_text)
        
        # 同时写入文件
        if self.params.log_to_file:
            try:
                self.log_file.write(log_text + '\n')
                self.log_file.flush()
            except Exception as e:
                print(f"写入日志文件失败: {e}")

    def safe_get_value(self, indicator, index=0):
        try:
            return float(indicator[index]) if not math.isnan(indicator[index]) else 0.0
        except:
            return 0.0
            
    def safe_get_bool(self, indicator, index=0):
        try:
            value = indicator[index]
            return bool(value) and not math.isnan(value) if not math.isnan(value) else False
        except:
            return False

    def get_current_foundation_state(self, jichu_b_val, jichu_s_val):
        """计算当前的基础状态 基础状态 := IF(基础B,1,IF(基础S,-1,0))"""
        if jichu_b_val:
            return 1
        elif jichu_s_val:
            return -1
        else:
            return 0

    def get_last_signal_state(self, states_history):
        """获取上一次信号的基础状态值（只能是1或-1）"""
        if not states_history or len(states_history) < 2:
            return 0
            
        # 从后往前查找第一个非0状态
        for i in range(len(states_history) - 2, -1, -1):
            date, state = states_history[i]
            if state != 0:
                return state
        return 0

    def next(self):
        current_date = self.datas[0].datetime.date(0)
        current_price = self.data_close[0]
        current_low = self.data_low[0]
        
        # 获取所有指标值
        p3_val = self.safe_get_value(self.p3)
        x_7_val = self.safe_get_value(self.x_7)
        x_8_val = self.safe_get_value(self.x_8)
        x_9_val = self.safe_get_value(self.x_9)
        sma1_val = self.safe_get_value(self.sma_layer1)
        sma2_val = self.safe_get_value(self.sma_layer2)
        
        # 获取条件布尔值
        hl_val = self.safe_get_bool(self.hl)
        ll_val = self.safe_get_bool(self.ll)
        zzd_val = self.safe_get_bool(self.zzd)
        lxyz_val = self.safe_get_bool(self.lxyz)
        jichu_b_val = self.safe_get_bool(self.jichu_b)
        jichu_s_val = self.safe_get_bool(self.jichu_s)
        
        # 存储基础状态历史
        if not hasattr(self, 'foundation_states'):
            self.foundation_states = []
            
        # 计算当前的基础状态
        current_foundation_state = self.get_current_foundation_state(jichu_b_val, jichu_s_val)
        
        # 保存当前的基础状态
        self.foundation_states.append((current_date, current_foundation_state))
        if len(self.foundation_states) > 100:
            self.foundation_states.pop(0)
            
        # 获取上一次信号的基础状态值
        last_signal_state = self.get_last_signal_state(self.foundation_states)
        
        # 检查有效信号
        # 有效B : 基础B AND REF(前次基础状态B,1) = -1  
        # 有效S : 基础S AND REF(前次基础状态B,1) = 1
        valid_buy = jichu_b_val and last_signal_state == -1  # 上一次是卖出信号
        valid_sell = jichu_s_val and last_signal_state == 1   # 上一次是买入信号
        
        # 完整的每日输出
        if self.params.daily_output:
            self.log(f'Price: {current_price:.4f}, Low: {current_low:.4f}')
            self.log(f'P3(EMA43): {p3_val:.4f}, X_7: {x_7_val:.4f}')
            self.log(f'SMA1: {sma1_val:.4f}, SMA2: {sma2_val:.4f}, X_8: {x_8_val:.4f}, X_9: {x_9_val:.4f}')
            self.log(f'HL(X_8>=X_9): {hl_val}, LL(X_8<=X_9): {ll_val}')
            self.log(f'ZZD(HL AND REF(LL,1)): {zzd_val}')
            self.log(f'LXYZ(COUNT(LL,10)=9): {lxyz_val}')
            self.log(f'C>P3: {current_price > p3_val}')
            self.log(f'基础B(ZZD AND LXYZ AND C>P3): {jichu_b_val}')
            self.log(f'基础S(CROSS(X_9,L)): {jichu_s_val}')
            self.log(f'当前基础状态: {current_foundation_state}')
            self.log(f'上一次信号状态: {last_signal_state} (1=买入,-1=卖出)')
            self.log(f'有效买入(需要上次为卖出): {valid_buy}')
            self.log(f'有效卖出(需要上次为买入): {valid_sell}')
            self.log(f'持仓数量: {self.position.size}')
            if len(self.foundation_states) >= 10:
                recent_states = [(str(date), state) for date, state in self.foundation_states[-10:]]
                self.log(f'最近基础状态历史: {recent_states}')
            self.log('-' * 80)
        
        # 买入信号处理
        if valid_buy:
            self.log(f'*** TDX买入信号触发 *** 有效B = 基础B AND 上次信号为卖出(-1)')
            
            # 存入数据库
            try:
                result = {
                    'strategy': 'TDXBSStrategyBuy',
                    'stock_code': self.datas[0]._name,
                    'op_act': 'buy_signal',
                    'op_price': float(x_9_val),#float(current_price),
                    'op_num': int(self.position.size),
                    'op_comm': 0.0,
                    'op_dt': current_date.isoformat(),
                    'reason': 'TDX_Buy_Signal_YouxiaoB',
                    'jichu_b': int(jichu_b_val),
                    'jichu_s': int(jichu_s_val),
                    'current_state': int(current_foundation_state),
                    'last_signal_state': int(last_signal_state),
                    'p3': float(p3_val),
                    'x_7': float(x_7_val),
                    'x_8': float(x_8_val),
                    'x_9': float(x_9_val),
                    'close_price': float(current_price),
                    'low_price': float(current_low),
                }
                DbOprate.save_result_to_mysql_all(result)
                self.log(f'*** TDX买入信号已存入数据库 ***')
                self.log(f'数据库存储数据: {result}')
            except Exception as e:
                self.log(f'*** 数据库存储失败 *** 错误: {str(e)}')
                
            if not self.position:
                self.buy()
                
        # 卖出信号处理
        elif valid_sell:
            self.log(f'*** TDX卖出信号触发 *** 有效S = 基础S AND 上次信号为买入(1)')
            
            # 存入数据库
            try:
                result = {
                    'strategy': 'TDXBSStrategySell',
                    'stock_code': self.datas[0]._name,
                    'op_act': 'sell_signal',
                    'op_price': float(x_9_val),#float(current_price),
                    'op_num': int(self.position.size),
                    'op_comm': 0.0,
                    'op_dt': current_date.isoformat(),
                    'reason': 'TDX_Sell_Signal_YouxiaoS',
                    'jichu_b': int(jichu_b_val),
                    'jichu_s': int(jichu_s_val),
                    'current_state': int(current_foundation_state),
                    'last_signal_state': int(last_signal_state),
                    'p3': float(p3_val),
                    'x_7': float(x_7_val),
                    'x_8': float(x_8_val),
                    'x_9': float(x_9_val),
                    'close_price': float(current_price),
                    'low_price': float(current_low),
                }
                DbOprate.save_result_to_mysql_all(result)
                self.log(f'*** TDX卖出信号已存入数据库 ***')
                self.log(f'数据库存储数据: {result}')
            except Exception as e:
                self.log(f'*** 数据库存储失败 *** 错误: {str(e)}')
                
            if self.position:
                self.sell()

    def stop(self):
        # 关闭日志文件
        if self.params.log_to_file and hasattr(self, 'log_file'):
            try:
                self.log_file.close()
                print(f"日志文件已保存")
            except Exception as e:
                print(f"关闭日志文件失败: {e}")

    def notify_order(self, order):
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'买入订单完成: 价格 {order.executed.price:.2f}, 数量: {order.executed.size}')
            elif order.issell():
                self.log(f'卖出订单完成: 价格 {order.executed.price:.2f}, 数量: {order.executed.size}')
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f'订单失败: 状态 {order.Status[order.status]}')
